from util.embeddings import embedding_filter
from util.entity import Article
from util.es import query_data


def es_recall(queries: list[str]):
    ds = query_data(index_name="law1", search_terms=queries, fields=["title", "content"])
    return ds


def recall(queries: list[str]) -> list[dict]:
    es_result = es_recall(queries)
    es_result = embedding_filter(queries, es_result, 0.5)
    es_result = [
        {"recall_type": "es",
         "result": Article(name=x["title"], content=x["content"], score=x["score"], file_id=x["file_id"],
                           permission_type=x.get("permission_type", 1), permission_set=x.get("permission_set", []))}
        for x in es_result]
    print("recall result: {}".format(es_result))
    return es_result
